use std::{
    collections::HashMap,
    net::SocketAddr,
    pin::Pin,
    sync::{Arc, RwLock},
    time::Duration,
};

use futures::{pin_mut, select_biased, AsyncRead, AsyncWrite, FutureExt, Stream, StreamExt};
use thiserror::Error as ErrorTrait;
use tokio::{
    net::UdpSocket,
    spawn,
    sync::{mpsc, oneshot},
    time::sleep,
};
use tracing::{error, warn};

use super::{Endpoint, OutgoingBiStreamsTrait};

mod secret;

#[derive(Debug, ErrorTrait)]
pub enum Error {
    #[error("IO错误")]
    Io(#[from] std::io::Error),
    #[error("连接错误")]
    Connect(#[from] quinn::ConnectError),
    #[error("链路错误")]
    Connection(#[from] quinn::ConnectionError),
    #[error("端点错误")]
    Endpoint(#[from] quinn::EndpointError),
    #[error("超时错误")]
    Timeout,
    #[error("重复的id")]
    DuplicateId,
    #[error("认证错误")]
    Certificate,
    #[error("监听停止")]
    StopListening,
    #[error("逻辑错误")]
    Logic,
}

pub struct QuicEndpointBuilder {
    ctx: IncomingContextSync,
    builder: quinn::EndpointBuilder,
}

impl QuicEndpointBuilder {
    pub fn new() -> Self {
        let ctx = Arc::new(RwLock::new(IncomingContext {
            connecting_ctxs: HashMap::new(),
        }));

        let mut config = quinn::ServerConfig::default();
        config.transport = std::sync::Arc::new({
            let mut transport_config = quinn::TransportConfig::default();
            transport_config.max_concurrent_uni_streams(0).unwrap();
            transport_config
        });
        Arc::get_mut(&mut config.crypto).map(|crypto| {
            crypto.cert_resolver = Arc::new(ResolvesServerCert { ctx: ctx.clone() });
            crypto
                .set_client_certificate_verifier(Arc::new(ClientCertVerifier { ctx: ctx.clone() }));
        });

        let mut builder = quinn::Endpoint::builder();
        builder.listen(config);

        Self { ctx, builder }
    }

    pub async fn bind(
        self,
        addr: &<QuicEndpoint as Endpoint>::EndpointAddr,
    ) -> Result<QuicEndpoint, Error> {
        let socket = UdpSocket::bind(addr.clone()).await?.into_std()?;
        self.builder
            .with_socket(socket)
            .map(|(ep, incoming)| {
                let incoming = Incoming::new(self.ctx, incoming);
                QuicEndpoint { ep, incoming }
            })
            .map_err(|e| e.into())
    }
}

struct ResolvesServerCert {
    ctx: IncomingContextSync,
}

impl rustls::ResolvesServerCert for ResolvesServerCert {
    fn resolve(&self, client_hello: rustls::ClientHello) -> Option<rustls::sign::CertifiedKey> {
        client_hello.server_name().and_then(|server_name| {
            self.ctx
                .write()
                .unwrap()
                .connecting_ctxs
                .get(server_name.into())
                .map(|c_ctx| c_ctx.server_cert.clone())
        })
    }
}

struct ClientCertVerifier {
    ctx: IncomingContextSync,
}

impl rustls::ClientCertVerifier for ClientCertVerifier {
    fn client_auth_root_subjects(
        &self,
        _sni: Option<&webpki::DNSName>,
    ) -> Option<rustls::DistinguishedNames> {
        Some(rustls::DistinguishedNames::new())
    }

    fn verify_client_cert(
        &self,
        presented_certs: &[rustls::Certificate],
        sni: Option<&webpki::DNSName>,
    ) -> std::result::Result<rustls::ClientCertVerified, rustls::TLSError> {
        sni.ok_or(rustls::TLSError::PeerIncompatibleError("No sni".into()))
            .and_then(|server_name| {
                let ctx = self.ctx.write().unwrap();

                ctx.connecting_ctxs
                    .get(server_name.as_ref().into())
                    .ok_or(rustls::TLSError::NoCertificatesPresented)
                    .and_then(|c_ctx| {
                        c_ctx
                            .client_cert_verifier
                            .verify_client_cert(presented_certs, sni)
                    })
            })
    }
}

pub struct QuicEndpoint {
    ep: quinn::Endpoint,
    incoming: Incoming,
}

#[derive(Clone)]
pub struct DomainIdentity {
    name: String,
    ca_der: Vec<u8>,
}

impl DomainIdentity {
    pub fn new(name: &str, ca_der: &[u8]) -> Self {
        Self {
            name: name.to_string(),
            ca_der: ca_der.to_vec(),
        }
    }
}

#[derive(Clone)]
pub struct EndpointIdentity {
    cert_der: Vec<u8>,
    key_der: Vec<u8>,
}

impl EndpointIdentity {
    pub fn new(cert_der: &[u8], key_der: &[u8]) -> Self {
        Self {
            cert_der: cert_der.to_vec(),
            key_der: key_der.to_vec(),
        }
    }
}

pub struct DomainIncoming {
    incoming: Incoming,
    name: String,
    rx: mpsc::Receiver<(OutgoingBiStreams, IncomingBiStreams)>,
}

impl futures::Stream for DomainIncoming {
    type Item = (OutgoingBiStreams, IncomingBiStreams);

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        self.as_mut().rx.poll_recv(cx)
    }
}

impl Drop for DomainIncoming {
    fn drop(&mut self) {
        self.incoming.stop_listen_for_domain(self.name.as_str());
    }
}

#[derive(Clone)]
struct Incoming {
    _handle: Arc<oneshot::Sender<()>>,
    ctx: IncomingContextSync,
}

type IncomingContextSync = Arc<RwLock<IncomingContext>>;

struct IncomingContext {
    connecting_ctxs: HashMap<String, ConnectingContext>,
}

struct ConnectingContext {
    server_cert: rustls::sign::CertifiedKey,
    client_cert_verifier: Arc<dyn rustls::ClientCertVerifier>,
    notifier: Option<ConnectingNotifier>,
}

enum ConnectingNotifier {
    Domain(mpsc::Sender<(OutgoingBiStreams, IncomingBiStreams)>),
    P2P(oneshot::Sender<(OutgoingBiStreams, IncomingBiStreams)>),
}

impl Incoming {
    fn new(ctx: IncomingContextSync, mut incoming: quinn::Incoming) -> Self {
        let (h_tx, h_rx) = oneshot::channel();
        let ctx1 = ctx.clone();

        spawn(async move {
            let rx_fut = h_rx.fuse();
            pin_mut!(rx_fut);

            loop {
                let connecting_fut = incoming.next().fuse();
                pin_mut!(connecting_fut);

                select_biased! {
                    _ = rx_fut => {
                        break;
                    }

                    connecting = connecting_fut => {
                        if let Some(connecting) = connecting {
                            match connecting.await {
                                Ok(conn) => {
                                    if let Some(handshake_data) = conn.connection.handshake_data() {
                                        if let Some(server_name) = handshake_data.server_name {
                                            let notifier = {
                                                let mut ctx = ctx.write().unwrap();
                                                ctx.connecting_ctxs.get_mut(&server_name)
                                                .and_then(|c_ctx| {
                                                    c_ctx.notifier.take()
                                                })
                                            };

                                            if let Some(notifier) = notifier {
                                                let split = split_conn(conn);

                                                match notifier {
                                                    ConnectingNotifier::Domain(tx) => {
                                                        if let Ok(_) = tx.send(split).await {
                                                            let mut ctx = ctx.write().unwrap();
                                                            ctx.connecting_ctxs.get_mut(&server_name)
                                                            .map(|c_ctx| {
                                                                c_ctx.notifier = Some(ConnectingNotifier::Domain(tx));
                                                            });
                                                        }
                                                    }

                                                    ConnectingNotifier::P2P(tx) => {
                                                        let mut ctx = ctx.write().unwrap();
                                                        ctx.connecting_ctxs.remove(&server_name);

                                                        tx.send(split).ok();
                                                    }
                                                }
                                            }
                                        } else {
                                            warn!("Duplicate connecting");
                                        }
                                    } else {
                                        error!("Logic error");
                                    }
                                }

                                Err(e) => {
                                    warn!("{}", e);
                                }
                            }
                        } else {
                            break;
                        }
                    }
                }
            }

            ctx.write().unwrap().connecting_ctxs.clear();
        });

        Self {
            _handle: Arc::new(h_tx),
            ctx: ctx1,
        }
    }

    fn listen_for_domain(
        &self,
        name: &str,
        ca_der: &[u8],
        cert_der: &[u8],
        key_der: &[u8],
    ) -> Result<mpsc::Receiver<(OutgoingBiStreams, IncomingBiStreams)>, Error> {
        let signing_key = rustls::sign::any_supported_type(&rustls::PrivateKey(key_der.to_vec()))
            .map_err(|_| Error::Certificate)?;
        let server_cert = rustls::sign::CertifiedKey {
            cert: vec![rustls::Certificate(cert_der.into())],
            key: Arc::new(signing_key),
            ocsp: None,
            sct_list: None,
        };

        let mut root_cert_store = rustls::RootCertStore::empty();
        root_cert_store
            .add(&rustls::Certificate(ca_der.into()))
            .map_err(|_| Error::Certificate)?;
        let client_cert_verifier = rustls::AllowAnyAuthenticatedClient::new(root_cert_store);

        let (tx, rx) = mpsc::channel(2);

        let connecting_ctx = ConnectingContext {
            server_cert,
            client_cert_verifier,
            notifier: Some(ConnectingNotifier::Domain(tx)),
        };

        self.ctx
            .write()
            .unwrap()
            .connecting_ctxs
            .insert(get_server_name_for_domain(name), connecting_ctx);

        Ok(rx)
    }

    fn stop_listen_for_domain(&self, name: &str) {
        let mut ctx = self.ctx.write().unwrap();
        ctx.connecting_ctxs
            .remove(&get_server_name_for_domain(name));
    }

    fn listen_for_p2p(
        &self,
        identity: &str,
        local_key: &secret::PrivateKey,
        peer_key: &secret::PublicKey,
    ) -> Result<oneshot::Receiver<(OutgoingBiStreams, IncomingBiStreams)>, Error> {
        todo!()
    }

    fn stop_listen_for_p2p(&self, identity: &str) {
        let mut ctx = self.ctx.write().unwrap();
        ctx.connecting_ctxs
            .remove(&get_server_name_for_p2p(identity));
    }
}

fn get_server_name_for_domain(name: &str) -> String {
    "d".to_string() + name
}

fn get_server_name_for_p2p(identity: &str) -> String {
    "p".to_string() + identity
}

fn split_conn(conn: quinn::NewConnection) -> (OutgoingBiStreams, IncomingBiStreams) {
    (
        OutgoingBiStreams {
            inner: conn.connection,
        },
        IncomingBiStreams {
            inner: conn.bi_streams,
        },
    )
}

pub struct OutgoingBiStreams {
    inner: quinn::Connection,
}

impl OutgoingBiStreamsTrait for OutgoingBiStreams {
    type Error = Error;
    type SendStream = SendStream;
    type RecvStream = RecvStream;

    async fn open_bi(&self) -> Result<(SendStream, RecvStream), Error> {
        self.inner
            .open_bi()
            .await
            .map(|(tx, rx)| (SendStream { inner: tx }, RecvStream { inner: rx }))
            .map_err(|e| e.into())
    }
}

pub struct IncomingBiStreams {
    inner: quinn::IncomingBiStreams,
}

impl Stream for IncomingBiStreams {
    type Item = Result<(SendStream, RecvStream), Error>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        quinn::IncomingBiStreams::poll_next(Pin::new(&mut self.inner), cx).map(|r| {
            r.map(|r| {
                r.map(|(tx, rx)| (SendStream { inner: tx }, RecvStream { inner: rx }))
                    .map_err(|e| e.into())
            })
        })
    }
}

pub struct SendStream {
    inner: quinn::SendStream,
}

impl AsyncWrite for SendStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        quinn::SendStream::poll_write(Pin::new(&mut self.inner), cx, buf)
    }

    fn poll_flush(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        quinn::SendStream::poll_flush(Pin::new(&mut self.inner), cx)
    }

    fn poll_close(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        quinn::SendStream::poll_close(Pin::new(&mut self.inner), cx)
    }
}

pub struct RecvStream {
    inner: quinn::RecvStream,
}

impl AsyncRead for RecvStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut [u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        quinn::RecvStream::poll_read(Pin::new(&mut self.inner), cx, buf)
    }
}

impl Endpoint for QuicEndpoint {
    type Error = Error;
    type DomainIdentity = DomainIdentity;
    type EndpointIdentity = EndpointIdentity;
    type EndpointAddr = SocketAddr;
    type PrivateKey = secret::PrivateKey;
    type PublicKey = secret::PublicKey;
    type DomainIncoming = DomainIncoming;
    type OutgoingBiStreams = OutgoingBiStreams;
    type IncomingBiStreams = IncomingBiStreams;
    type SendStream = SendStream;
    type RecvStream = RecvStream;

    async fn listen_for_domain(
        &self,
        domain_identity: &Self::DomainIdentity,
        endpoint_identity: &Self::EndpointIdentity,
    ) -> Result<Self::DomainIncoming, Self::Error> {
        self.incoming
            .listen_for_domain(
                domain_identity.name.as_str(),
                &domain_identity.ca_der[..],
                &endpoint_identity.cert_der[..],
                &endpoint_identity.key_der[..],
            )
            .map(|rx| DomainIncoming {
                incoming: self.incoming.clone(),
                name: domain_identity.name.clone(),
                rx,
            })
    }

    async fn connect_for_domain(
        &self,
        peer_addr: &Self::EndpointAddr,
        domain_identity: &Self::DomainIdentity,
        endpoint_identity: &Self::EndpointIdentity,
    ) -> Result<(Self::OutgoingBiStreams, Self::IncomingBiStreams), Self::Error> {
        let mut config = quinn::ClientConfig::default();
        config.transport = std::sync::Arc::new({
            let mut transport_config = quinn::TransportConfig::default();
            transport_config.max_concurrent_uni_streams(0).unwrap();
            transport_config
        });
        config
            .add_certificate_authority(
                quinn::Certificate::from_der(domain_identity.ca_der.as_slice())
                    .map_err(|_| Error::Certificate)?,
            )
            .map_err(|_| Error::Certificate)?;
        Arc::get_mut(&mut config.crypto)
            .ok_or(Error::Logic)
            .and_then(|crypto| {
                crypto
                    .set_single_client_cert(
                        vec![rustls::Certificate(endpoint_identity.cert_der.clone())],
                        rustls::PrivateKey(endpoint_identity.key_der.clone()),
                    )
                    .map_err(|_| Error::Certificate)
            })?;

        let server_name = get_server_name_for_domain(domain_identity.name.as_str());
        let connecting = self
            .ep
            .connect_with(config, peer_addr, server_name.as_str())?;
        connecting
            .await
            .map(|conn| split_conn(conn))
            .map_err(|e| e.into())
    }

    async fn connect_for_p2p(
        &self,
        peer_addr: &Self::EndpointAddr,
        identity: &str,
        local_key: &Self::PrivateKey,
        peer_key: &Self::PublicKey,
        active: bool,
    ) -> Result<(Self::OutgoingBiStreams, Self::IncomingBiStreams), Self::Error> {
        if !active {
            // TODO: NAT
            let send_zero_raw = async {
                loop {
                    self.ep.send_zero_raw_packet(peer_addr).await.ok();
                    sleep(Duration::from_millis(500)).await;
                }
            };

            let rx = self
                .incoming
                .listen_for_p2p(identity, local_key, peer_key)?;

            struct Cleaner<'a, 'b> {
                incoming: &'a Incoming,
                identity: &'b str,
            }

            impl<'a, 'b> Drop for Cleaner<'_, '_> {
                fn drop(&mut self) {
                    self.incoming.stop_listen_for_p2p(self.identity);
                }
            }

            let _c = Cleaner {
                incoming: &self.incoming,
                identity,
            };

            select_biased! {
                r = rx.fuse() => {
                    r.map_err(|_| Error::StopListening)
                }

                _ = send_zero_raw.fuse() => {
                    Err(Error::Timeout)
                }
            }
        } else {
            let mut config = quinn::ClientConfig::default();
            config.transport = std::sync::Arc::new({
                let mut transport_config = quinn::TransportConfig::default();
                transport_config.max_concurrent_uni_streams(0).unwrap();
                transport_config
            });
            config
                .add_certificate_authority(
                    quinn::Certificate::from_der(peer_key.as_ref())
                        .map_err(|_| Error::Certificate)?,
                )
                .map_err(|_| Error::Certificate)?;
            Arc::get_mut(&mut config.crypto)
                .ok_or(Error::Logic)
                .and_then(|crypto| {
                    let (cert_der, key_der) = local_key.create_signed();
                    crypto
                        .set_single_client_cert(
                            vec![rustls::Certificate(cert_der)],
                            rustls::PrivateKey(key_der),
                        )
                        .map_err(|_| Error::Certificate)
                })?;

            todo!()
        }
    }

    async fn wait_idle(&self) {
        self.ep.wait_idle().await;
    }
}
